import mindspore as ms
import mindspore.nn as nn
from mindvision.classification.dataset import Mnist

from util.callback import AccuracyMonitor


def main():
    batch_size, lr, num_epochs = 256, 0.1, 10

    dataset_train = Mnist(path='/shareData/mindspore-dataset/Mnist', split="train",
                          batch_size=batch_size, repeat_num=1, shuffle=True,
                          resize=224, download=False).run()
    dataset_test = Mnist(path='/shareData/mindspore-dataset/Mnist', split="test",
                         batch_size=batch_size, repeat_num=1, shuffle=True,
                         resize=224, download=False).run()

    AlexNet = nn.SequentialCell(
        nn.Conv2d(1, 6, 5, 1),
        nn.Sigmoid(),
        nn.MaxPool2d(2, 2),
        nn.Conv2d(6, 16, 5, 1),
        nn.Sigmoid(),
        nn.MaxPool2d(2, 2),
        nn.Flatten(),
        nn.Dense(120, 84),
        nn.ReLU(),
        nn.Dense(84, 10),
        nn.Softmax()
    )

    loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
    opti = nn.SGD(AlexNet.trainable_params(), learning_rate=lr)

    model = ms.Model(AlexNet, loss_fn=loss, optimizer=opti, metrics={'acc'})

    model.train(num_epochs, dataset_train, callbacks=[AccuracyMonitor(dataset_test)])


if __name__ == '__main__':
    main()
